{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numbers\n",
    "from PIL import Image, ImageMath\n",
    "import os\n",
    "import os.path\n",
    "import numpy as np\n",
    "import struct\n",
    "import math\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('TkAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "import time\n",
    "import itertools\n",
    "\n",
    "import pcl\n",
    "import timeit\n",
    "import png\n",
    "\n",
    "import multiprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def axisEqual3D(ax):\n",
    "    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])\n",
    "    sz = extents[:, 1] - extents[:, 0]\n",
    "    centers = np.mean(extents, axis=1)\n",
    "    maxsize = max(abs(sz))\n",
    "    r = maxsize / 2\n",
    "    for ctr, dim in zip(centers, 'xyz'):\n",
    "        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)\n",
    "\n",
    "\n",
    "def plot_pc(pc_np, z_cutoff=1000, birds_view=False, color='height', size=0.3, ax=None, cmap=cm.jet, is_equal_axes=True):\n",
    "    # remove large z points\n",
    "    valid_index = pc_np[:, 0] < z_cutoff\n",
    "    pc_np = pc_np[valid_index, :]\n",
    "\n",
    "    if ax is None:\n",
    "        fig = plt.figure(figsize=(9, 9))\n",
    "        ax = Axes3D(fig)\n",
    "    if type(color)==str and color == 'height':\n",
    "        c = pc_np[:, 2]\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=c, cmap=cmap, edgecolors='none')\n",
    "    elif type(color)==str and color == 'reflectance':\n",
    "        assert False\n",
    "    elif type(color) == np.ndarray:\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=color, cmap=cmap, edgecolors='none')\n",
    "    else:\n",
    "        ax.scatter(pc_np[:, 0], pc_np[:, 1], pc_np[:, 2], s=size, c=color, edgecolors='none')\n",
    "\n",
    "    if is_equal_axes:\n",
    "        axisEqual3D(ax)\n",
    "    if True == birds_view:\n",
    "        ax.view_init(elev=0, azim=-90)\n",
    "    else:\n",
    "        ax.view_init(elev=-45, azim=-90)\n",
    "    # ax.invert_yaxis()\n",
    "\n",
    "    return ax\n",
    "\n",
    "\n",
    "def read_txt_as_list(txt_file):\n",
    "    with open(txt_file, 'r') as f:\n",
    "        subset = [x.rstrip() for x in f.readlines()]\n",
    "    return subset\n",
    "\n",
    "\n",
    "def read_camera_intrinsics(txt_file):\n",
    "    return np.loadtxt(txt_file)\n",
    "\n",
    "def img2pc(depth, camera_matrix):\n",
    "    \"\"\"Transform a depth image into a point cloud with one point for each\n",
    "    pixel in the image, using the camera transform for a camera\n",
    "    centred at cx, cy with field of view fx, fy.\n",
    "\n",
    "    depth is a 2-D ndarray with shape (rows, cols) containing\n",
    "    depths from 1 to 254 inclusive. The result is a 3-D array with\n",
    "    shape (rows, cols, 3). Pixels with invalid depth in the input have\n",
    "    NaN for the z-coordinate in the result.\n",
    "\n",
    "    \"\"\"\n",
    "    fx = camera_matrix[0, 0]\n",
    "    fy = camera_matrix[1, 1]\n",
    "    cx = camera_matrix[0, 2]\n",
    "    cy = camera_matrix[1, 2]\n",
    "    \n",
    "    rows, cols = depth.shape\n",
    "    c, r = np.meshgrid(np.arange(cols), np.arange(rows), sparse=True)\n",
    "    valid = (depth > 1)\n",
    "    z = np.where(valid, depth / 1000.0, 0)\n",
    "    x = np.where(valid, z * (c - cx) / fx, 0)\n",
    "    y = np.where(valid, z * (r - cy) / fy, 0)\n",
    "    \n",
    "    pc_HxWx3 = np.dstack((x, y, z))\n",
    "    pc = pc_HxWx3.reshape(rows*cols, 3)\n",
    "    valid_idx = pc[:, 2] > 1e-3\n",
    "    pc = pc[valid_idx, :]\n",
    "    \n",
    "    return pc\n",
    "\n",
    "\n",
    "# get surface normal\n",
    "def Surface_normals(cloud):\n",
    "    ne = cloud.make_NormalEstimation()\n",
    "    tree = cloud.make_kdtree()\n",
    "    ne.set_SearchMethod(tree)\n",
    "#     ne.set_RadiusSearch(2)\n",
    "    ne.set_KSearch(9)\n",
    "    cloud_normals = ne.compute()\n",
    "    return cloud_normals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# single thread\n",
    "dataset_root = '/data/datasets/3DMatch_train'\n",
    "output_root = '/data/datasets/3DMatch_npy'\n",
    "downsample_size = 16384\n",
    "\n",
    "train_list = read_txt_as_list(os.path.join(dataset_root, 'training_list.txt'))\n",
    "\n",
    "for folder in train_list:\n",
    "    # load camera intrinsics\n",
    "    camera_matrix = read_camera_intrinsics(os.path.join(dataset_root, folder, 'camera-intrinsics.txt'))\n",
    "    \n",
    "    # find subfolder\n",
    "    things_in_folder = os.listdir(os.path.join(dataset_root, folder));\n",
    "    for subfolder in things_in_folder:\n",
    "        if not os.path.isdir(os.path.join(dataset_root, folder, subfolder)):\n",
    "            continue\n",
    "        # print info\n",
    "        print(os.path.join(dataset_root, folder, subfolder))\n",
    "        \n",
    "        # find files\n",
    "        things_in_subfolder = os.listdir(os.path.join(dataset_root, folder, subfolder))\n",
    "        for file in things_in_subfolder:\n",
    "            if 'depth.png' in file:\n",
    "                png_file = os.path.join(dataset_root, folder, subfolder, file)\n",
    "                npy_folder = os.path.join(output_root, folder, subfolder)\n",
    "                if not os.path.isdir(npy_folder):\n",
    "                    os.makedirs(npy_folder)\n",
    "                npy_file = os.path.join(npy_folder, file[0:-4]+'.npy')\n",
    "                \n",
    "                png_reader = png.Reader(png_file)\n",
    "                png_data = png_reader.read()\n",
    "                png_np = np.vstack(list(map(np.uint16, png_data[2])))\n",
    "                \n",
    "                pc_np = img2pc(png_np, camera_matrix)\n",
    "                # downsample\n",
    "                if pc_np.shape[0] > downsample_size:\n",
    "                    choice_idx = np.random.choice(pc_np.shape[0], downsample_size, replace=False)\n",
    "                    pc_np = pc_np[choice_idx, :]\n",
    "                    \n",
    "                # compute surface normal\n",
    "                cloud = pcl.PointCloud(pc_np.astype(np.float32))    \n",
    "                sn = Surface_normals(cloud)\n",
    "                sn_np = np.asarray(sn.to_array(), dtype=np.float32)  # Nx4, nx,ny,nz,curvature\n",
    "\n",
    "                output_np = np.concatenate((pc_np.astype(np.float32), sn_np), axis=1)  # Nx7\n",
    "                np.save(npy_file, output_np)\n",
    "                \n",
    "                # visualize\n",
    "                # plot_pc(pc_np)\n",
    "                # plt.show()\n",
    "                # break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-06\n",
      "/data/datasets/3DMatch_train/sun3d-hotel_umd-maryland_hotel1/seq-01\n",
      "/data/datasets/3DMatch_train/sun3d-home_md-home_md_scan9_2012_sep_30/seq-01\n",
      "/data/datasets/3DMatch_train/sun3d-mit_76_studyroom-76-1studyroom2/seq-01\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-04\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-03\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-11\n",
      "/data/datasets/3DMatch_train/sun3d-mit_lab_hj-lab_hj_tea_nov_2_2012_scan1_erika/seq-01\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-05\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-13\n",
      "/data/datasets/3DMatch_train/sun3d-hotel_umd-maryland_hotel3/seq-01\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-07\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-02\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-14\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-12\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-01\n",
      "/data/datasets/3DMatch_train/7-scenes-redkitchen/seq-08\n",
      "/data/datasets/3DMatch_train/sun3d-home_at-home_at_scan1_2013_jan_1/seq-01\n",
      "/data/datasets/3DMatch_train/sun3d-hotel_uc-scan3/seq-01\n"
     ]
    }
   ],
   "source": [
    "# multithread\n",
    "total_thread_num = 4\n",
    "def save_npy(thread_idx):\n",
    "    dataset_root = '/data/datasets/3DMatch_train'\n",
    "    output_root = '/data/datasets/3DMatch_npy'\n",
    "    downsample_size = 16384\n",
    "\n",
    "    train_list = read_txt_as_list(os.path.join(dataset_root, 'testing_list.txt'))\n",
    "    \n",
    "    chunk_size = math.floor(len(train_list) / total_thread_num)\n",
    "    if thread_idx != total_thread_num-1:\n",
    "        begin_idx = thread_idx * chunk_size\n",
    "        end_idx = (thread_idx+1) * chunk_size\n",
    "    else:\n",
    "        begin_idx = thread_idx * chunk_size\n",
    "        end_idx = len(train_list)\n",
    "\n",
    "    for folder in train_list[begin_idx:end_idx]:\n",
    "        # load camera intrinsics\n",
    "        camera_matrix = read_camera_intrinsics(os.path.join(dataset_root, folder, 'camera-intrinsics.txt'))\n",
    "\n",
    "        # find subfolder\n",
    "        things_in_folder = os.listdir(os.path.join(dataset_root, folder));\n",
    "        for subfolder in things_in_folder:\n",
    "            if not os.path.isdir(os.path.join(dataset_root, folder, subfolder)):\n",
    "                continue\n",
    "            # print info\n",
    "            print(os.path.join(dataset_root, folder, subfolder))\n",
    "\n",
    "            # find files\n",
    "            things_in_subfolder = os.listdir(os.path.join(dataset_root, folder, subfolder))\n",
    "            for file in things_in_subfolder:\n",
    "                if 'depth.png' in file:\n",
    "                    png_file = os.path.join(dataset_root, folder, subfolder, file)\n",
    "                    npy_folder = os.path.join(output_root, folder, subfolder)\n",
    "                    if not os.path.isdir(npy_folder):\n",
    "                        os.makedirs(npy_folder)\n",
    "                    npy_file = os.path.join(npy_folder, file[0:-4]+'.npy')\n",
    "\n",
    "                    png_reader = png.Reader(png_file)\n",
    "                    png_data = png_reader.read()\n",
    "                    png_np = np.vstack(list(map(np.uint16, png_data[2])))\n",
    "\n",
    "                    pc_np = img2pc(png_np, camera_matrix)\n",
    "                    # downsample\n",
    "                    if pc_np.shape[0] > downsample_size:\n",
    "                        choice_idx = np.random.choice(pc_np.shape[0], downsample_size, replace=False)\n",
    "                        pc_np = pc_np[choice_idx, :]\n",
    "\n",
    "                    # compute surface normal\n",
    "                    cloud = pcl.PointCloud(pc_np.astype(np.float32))    \n",
    "                    sn = Surface_normals(cloud)\n",
    "                    sn_np = np.asarray(sn.to_array(), dtype=np.float32)  # Nx4, nx,ny,nz,curvature\n",
    "\n",
    "                    output_np = np.concatenate((pc_np.astype(np.float32), sn_np), axis=1)  # Nx7\n",
    "                    np.save(npy_file, output_np)\n",
    "\n",
    "                    # visualize\n",
    "                    # plot_pc(pc_np)\n",
    "                    # plt.show()\n",
    "                    # break\n",
    "\n",
    "p_list = []\n",
    "for i in range(total_thread_num):\n",
    "    p_list.append(0)\n",
    "for i in range(total_thread_num):\n",
    "    p_list[i] = multiprocessing.Process(target=save_npy, args=(i, ))\n",
    "    p_list[i].start()\n",
    "    \n",
    "for i in range(total_thread_num):\n",
    "    p_list[i].join()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
